#-*- coding:utf-8 -*-
import numpy as np
import matplotlib.pyplot as mp
from partie1 import *

xa = 0.
ya = 0.
ma = 10.
mb = 1.
t0 = np.array([[0.]])
y0 = np.array([[1.],
               [0.],
               [0.],
               [2.]])

f_1 = lambda X,T : np.array([[X[2][0]],
                             [X[3][0]],
                             [ma*mb*(xa-X[0][0])/(((xa-X[0][0])**2 + (ya-X[1][0])**2)**(1.5))],
                             [ma*mb*(ya-X[1][0])/(((xa-X[0][0])**2 + (ya-X[1][0])**2)**(1.5))]])

def plot_sol_three_points(t0,y0,f):
    "Affiche la solution approximée sur l' intervalle [0 10]"
    mp.clf()
    nmax = 500
    tmax = 15.
    res = meth_n_step(y0, t0, nmax, tmax/nmax, f,step_kutta)[0]
    xb = []
    yb = []
    dxb = []
    dyb = []
    for i in range(len(res)):
        xb = xb + [res[i][0][0]]
        yb = yb + [res[i][1][0]]
        dxb = dxb + [res[i][2][0]]
        dyb = dyb + [res[i][3][0]]
    mp.plot(xb, yb, linewidth=1.0)
    mp.scatter(0.,0., c="yellow", s = 200)
    mp.scatter(xb[0],yb[0])
    mp.show()

plot_sol_three_points(t0,y0,f_1)

# Comportement erroné : lorsque B approche trop A, la résolution s'emballe et envoie B vers l'infini et au délà

xa = 0.
ya = 0.
ma = 1.
mb = 0.01
#xb(t): cos(t)
#yb(t): sin(t) car |AB|=1, trajectoire circulaire et période de 2pi

f_2 = lambda X,T : np.array([[X[2][0]],
                             [X[3][0]],
                             [(ma*(xa-X[0][0])/(((xa-X[0][0])**2 + (ya-X[1][0])**2)**(1.5))) + (mb*(np.cos(T[0][0])-X[0][0])/(((np.cos(T[0][0])-X[0][0])**2 + (np.sin(T[0][0])-X[1][0])**2)**(1.5)))],
                             [(ma*(ya-X[1][0])/(((xa-X[0][0])**2 + (ya-X[1][0])**2)**(1.5))) + (mb*(np.sin(T[0][0])-X[1][0])/(((np.cos(T[0][0])-X[0][0])**2 + (np.sin(T[0][0])-X[1][0])**2)**(1.5)))]])


# Nous avons essayé de réutiliser la force centrifuge du précédent TP pour nous ramené dans un repère
# héliocentrique en rotation avec la Terre... useless ?
# f_3 = lambda X,T : np.array([[X[2][0]],
#                              [X[3][0]],
#                              [(ma*(xa-X[0][0])/(((xa-X[0][0])**2 + (ya-X[1][0])**2)**(1.5))) + (mb*(np.cos(T[0][0])-X[0][0])/(((np.cos(T[0][0])-X[0][0])**2 + (np.sin(T[0][0])-X[1][0])**2)**(1.5))) + ma*(X[0][0] - ((ma*xa+mb*T[0][0])/(ma+mb)))],
#                              [(ma*(ya-X[1][0])/(((xa-X[0][0])**2 + (ya-X[1][0])**2)**(1.5))) + (mb*(np.sin(T[0][0])-X[1][0])/(((np.cos(T[0][0])-X[0][0])**2 + (np.sin(T[0][0])-X[1][0])**2)**(1.5))) + ma*(X[0][0] - ((ma*xa+mb*T[0][0])/(ma+mb)))]])


t0 = np.array([[0.]])
y0 = np.array([[0.95],[0.95],[-0.5],[0.5]])

def plot_sol_three_points_2(t0,y0,f):
    "Affiche le mouvement d'un satellite en présence de la Terre et du Soleil"
    mp.clf()
    nmax = 1000
    tmax = 7.
    res = meth_n_step(y0, t0, nmax, tmax/nmax, f,step_kutta)[0]
    temps = np.arange(0.,tmax*(1.+1./nmax),tmax/nmax)
    xc = []
    yc = []
    for i in range(len(res)):
        xc = xc + [res[i][0][0]]
        yc = yc + [res[i][1][0]]
    c = mp.plot(xc, yc, linewidth=1.0)
    xb = map(np.cos, np.arange(0, 2*np.pi+1, 0.1))
    yb = map(np.sin, np.arange(0, 2*np.pi+1, 0.1))
    b = mp.plot(xb, yb, linewidth=1.0)
    mp.scatter(0.,0., c="yellow", s = 200)
    mp.scatter(xc[0],yc[0])
    mp.scatter(1,0, c="green")
    mp.legend((c,b),("Trajectoire de C","Trajectoire de B"))
    mp.show()

plot_sol_three_points_2(t0,y0,f_2)


def plot_sol_three_points_3(t0,f):
    "affiche deux mouvements d'un satellite pour des conditions initiales proches"
    mp.clf()
    y1 = np.array([[-0.3],[0.6],[-0.5],[-0.5]])
    y2 = np.array([[-0.31],[0.61],[-0.5],[-0.5]])#pour deux conditions initiales très proches
    nmax = 1000
    tmax = 7.
    res1 = meth_n_step(y1, t0, nmax, tmax/nmax, f,step_kutta)[0]
    res2 = meth_n_step(y2, t0, nmax, tmax/nmax, f,step_kutta)[0]
    temps = np.arange(0.,tmax*(1.+1./nmax),tmax/nmax)
    xc1 = []
    yc1 = []
    xc2 = []
    yc2 = []
    for i in range(len(res1)):
        xc1 = xc1 + [res1[i][0][0]]
        yc1 = yc1 + [res1[i][1][0]]
        xc2 = xc2 + [res2[i][0][0]]
        yc2 = yc2 + [res2[i][1][0]]
    rxc1 = []
    ryc1 = []
    rxc2 = []
    ryc2 = []
    for i in range(len(xc1)):#rotation de la trajectoire
        rxc1 = rxc1 + [xc1[i]*np.cos(temps[i]) + yc1[i]*np.sin(temps[i])]
        ryc1 = ryc1 + [yc1[i]*np.cos(temps[i]) - xc1[i]*np.sin(temps[i])]
        rxc2 = rxc2 + [xc2[i]*np.cos(temps[i]) + yc2[i]*np.sin(temps[i])]
        ryc2 = ryc2 + [yc2[i]*np.cos(temps[i]) - xc2[i]*np.sin(temps[i])]
    c1 = mp.plot(rxc1, ryc1, linewidth=1.0)
    c2 = mp.plot(rxc2, ryc2, linewidth=1.0)
    xb = map(np.cos, np.arange(0, 2*np.pi+1, 0.1))
    yb = map(np.sin, np.arange(0, 2*np.pi+1, 0.1))
    b = mp.plot(xb, yb, linewidth=1.0)
    mp.scatter(0.,0., c="yellow", s = 200)
    mp.scatter(1.,0., c="red", s = 50)
    mp.scatter(xc1[0],yc1[0])
    mp.scatter(xc2[0],yc2[0], c="green")
    mp.scatter(rxc1[len(xc1)-1],ryc1[len(xc1)-1])
    mp.scatter(rxc2[len(xc2)-1],ryc2[len(xc2)-1], c="green")
    mp.legend((c1,c2,b),("Trajectoire de C1","Trajectoire de C2","Trajectoire de B"))
    mp.show()

plot_sol_three_points_3(t0,f_2)


def plot_sol_three_points_4(t0,f):
    "Points de Lagrange - début et recherche de l'équilibre stable"
    mp.clf()
    # ci
    theta = np.pi/3
    y0 = np.array([[np.cos(theta)],
                   [np.sin(theta)],
                   [-np.sin(theta)],
                   [np.cos(theta)]])#sommet du triangle équilatérale
    nmax = 500
    tmax = 20. * np.pi
    # resolution et mise sous forme X Y
    res = meth_n_step(y0, t0, nmax, tmax/nmax, f,step_kutta)[0]
    temps = np.arange(0.,tmax*(1.+1./nmax),tmax/nmax)
    xc = []
    yc = []
    for i in range(len(res)):
        xc = xc + [res[i][0][0]]
        yc = yc + [res[i][1][0]]
    # rotation
    rxc = []
    ryc = []
    for i in range(len(xc)):
        rxc = rxc + [xc[i]*np.cos(temps[i]) + yc[i]*np.sin(temps[i])]
        ryc = ryc + [yc[i]*np.cos(temps[i]) - xc[i]*np.sin(temps[i])]
    # plot
    c = mp.plot(xc, yc, linewidth=1.0)
    mp.scatter(0.,0., c="yellow", s = 200)
    mp.scatter(1.,0., c="red", s = 50)
    mp.scatter(xc[0],yc[0])
    mp.legend((c,b),("Trajectoire de C1","Trajectoire de B"))
    mp.show()
    mp.plot(rxc, ryc, linewidth=1.0)
    mp.scatter(0.,0., c="yellow", s = 200)
    mp.scatter(1.,0., c="red", s = 50)
    mp.scatter(xc[0],yc[0])
    mp.show()

plot_sol_three_points_4(t0,f_2)
